Skip to content

Feat: add BaseSessionService.flush() and drain it from the Runner (adk-python parity) - #818

Open
AmaadMartin wants to merge 2 commits into
mainfrom
feat/session-service-flush-hook
Open

Feat: add BaseSessionService.flush() and drain it from the Runner (adk-python parity)#818
AmaadMartin wants to merge 2 commits into
mainfrom
feat/session-service-flush-hook

Conversation

@AmaadMartin

@AmaadMartin AmaadMartin commented Aug 8, 2026

Copy link
Copy Markdown
Owner

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

  1. Link to an existing issue (if applicable):
    N/A
  2. Or, if no issue exists, describe the change:
    Problem: A session service that batches event writes has no defined point at which to drain its buffer. adk-js offers no hook, and the Runner never signals the end of an invocation to the session service. adk-python has this hook and adk-js does not, so this is a cross-language parity gap.

Solution: Add a concrete, default no-op flush() to BaseSessionService, and await it in the finally of Runner.runAsync, after the toolsets close. adk-python calls flush() from Runner.close() (src/google/adk/runners.py); adk-js has no Runner.close(), and the runAsync teardown is where adk-js already performs the toolset cleanup that adk-python performs in close(). Draining per invocation is a stricter contract than draining at shutdown, and it gives the hook a caller today. No shipped session service overrides flush(), so nothing observable changes: InMemorySessionService, DatabaseSessionService and VertexAiSessionService all write synchronously.

Parity notes: the method name, the no-op default and the position after appendEvent follow src/google/adk/sessions/base_session_service.py. The call site is the one deliberate deviation, explained above. The Python if self.session_service guard is dropped because Runner.sessionService is a required field here.

Collision check: I listed the 400 open PRs on this fork and read the diffs of every adjacent one. No PR adds flush. #711, #616 and #609 also create core/test/sessions/base_session_service_test.ts, so whichever merges second needs a trivial add/add resolution. #782 and #716 touch runner.ts in other methods.

Testing Plan

Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes.
Unit Tests:
[x] I have added or updated unit tests for my change.
[x] All unit tests pass locally.

npx vitest run --project unit:core core/test/runner/runner_test.ts core/test/sessions/base_session_service_test.ts
# Test Files  2 passed (2)   Tests  39 passed (39)
npm run lint          # clean
npm run format:check  # clean
bash scripts/check_license.sh  # clean

Six new cases in core/test/runner/runner_test.ts (new describe, no existing block edited) and two in the new core/test/sessions/base_session_service_test.ts. They pin: one drain per completed invocation, the drain after the last appendEvent, one drain per invocation rather than per runner, a drain on the error path, a drain before runEphemeral deletes the session, and a real buffering subclass whose buffer is empty and whose events are written when the invocation ends.

Coverage: both new source lines are covered. core/src/runner/runner.ts:457 records 36 hits and the flush() body in core/src/sessions/base_session_service.ts records 1, measured with --coverage.reporter=json over the two targeted files.

Proof the tests can fail. I deleted await this.sessionService.flush(); from runner.ts and re-ran; all six runner cases failed:

× drains the session service once per completed invocation
  → expected "flush" to be called 1 times, but got 0 times
× drains the session service after the last appendEvent
  → actual value must be number or bigint, received "undefined"
× drains the session service once per invocation, not once per runner
  → expected "flush" to be called 2 times, but got 0 times
× drains the session service when the invocation throws
  → expected "flush" to be called 1 times, but got 0 times
× drains the session service before runEphemeral deletes the session
  → expected "flush" to be called 1 times, but got 0 times
× drains a buffering session service by the end of the invocation
  → expected [ { …(9), …(1) }, { …(8), …(1) } ] to deeply equal []
Tests  6 failed | 31 passed (37)

I also moved the call from the finally to the top of runAsync, upstream of the agent loop. Only the ordering case failed, which shows it pins order and not merely presence:

× drains the session service after the last appendEvent
  → expected 12 to be greater than 14
Tests  1 failed | 35 passed (36)

Finally I changed the base body to throw new Error('x'). Both base cases failed:

× resolves to undefined by default
  → promise rejected "Error: x" instead of resolving
× is inherited by a shipped service and leaves its events untouched
  → promise rejected "Error: x" instead of resolving

Manual End-to-End (E2E) Tests:
Please provide instructions on how to manually test your changes, including any necessary setup or configuration.

There is no user-visible behaviour to exercise by hand, because no shipped service buffers. To confirm the extension point by hand, subclass InMemorySessionService with a flush() that increments a counter, drive it through runner.runEphemeral(...), and read the counter. It reads 1. The drains a buffering session service by the end of the invocation test is the automated form of this, and it uses a real Runner and a real subclass with no mocks.

Checklist

[x] I have read the CONTRIBUTING.md document.
[x] I have performed a self-review of my own code.
[x] I have commented my code, particularly in hard-to-understand areas.
[x] I have added tests that prove my fix is effective or that my feature works.
[x] New and existing unit tests pass locally with my changes.

CI note

The first run-tests (windows-latest) job failed on core/test/code_executors/unsafe_local_code_executor_test.ts > should execute shell code and return stdout with Test timed out in 5000ms. That test spawns a shell and touches neither the Runner nor a session service, so this change cannot reach it. I re-ran the job with no code change and it passed. All six checks are green.

Amaad Martin added 2 commits August 8, 2026 11:37
…Runner

A session service that buffers writes has no defined point at which to
drain. adk-python gives it one: a default no-op BaseSessionService.flush().
Add the same hook here and call it in the Runner's invocation teardown,
after the toolsets close.
Pin the four invariants: the Runner drains once per invocation, after the
last appendEvent, on the error path, and before runEphemeral deletes the
session. A real buffering subclass proves the hook works end to end.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant